import os
import jieba
from Tools import savefile, readfile,load_stopwords
import re

def corpus_segment(corpus_path, seg_path, stopwords_file_path):
    stopwords = load_stopwords(stopwords_file_path)
    catelist = os.listdir(corpus_path)  # 获取corpus_path下的所有子目录
    #print(catelist)
    '''
    其中子目录的名字就是类别名，例如：
    train_corpus/体育/21.txt中，'train_corpus/'是corpus_path，'体育'是catelist中的一个成员
    '''
    # 获取每个目录（类别）下所有的文件
    for mydir in catelist:
        class_path = corpus_path + mydir + "/"  # 拼出分类子目录的路径如：train_corpus/体育/
        seg_dir = seg_path + mydir +  "/" # 拼出分类分词目录的路径如：train_seg/体育/
        if not os.path.exists(seg_dir):  # 是否存在分词目录
            os.makedirs(seg_dir)        # 如果没有则创建该目录

        file_list = os.listdir(class_path)  # 获取未分词语料库中某一类别中的所有文本
        for file_path in file_list:  # 遍历类别目录下的所有文件
            fullname = class_path + file_path  # 拼出文件名全路径如：train_corpus/art/21.txt
            content = readfile(fullname)  # 读取文件内容


            # 将字节字符串转换为普通字符串
            content = content.decode('utf-8')
            # 使用正则表达式去除空白字符
            content = re.sub(r'\s+', '', content)
            # 删除数字和英文字符
            filter_pattern = re.compile(r'[a-zA-Z0-9]+')
            content = re.sub(filter_pattern, '', content)
            # 进行分词
            # content_seg = jieba.cut_for_search(content)
            content_seg = jieba.cut(content)
            # 去停用词
            content_seg = [word for word in content_seg if word not in stopwords]
            # 将分词结果转换为字符串
            content_seg_str = ' '.join(content_seg)
            # 将分词结果保存到文件
            savefile(seg_dir + file_path, content_seg_str.encode('utf-8'))

if __name__ == "__main__":

    stopwords_file = "./stopwords.txt"  # 停用词文件路径

    corpus_path = "./train_corpus/"     # 需要分词的语料库路径
    seg_path = "./train_corpus_seg/"    # 分词之后的语料库储存路径
    corpus_segment(corpus_path, seg_path, stopwords_file)
    print("训练语料分词结束！！！")

    corpus_path = "./test_corpus/"  # 需要分词的语料库路径
    seg_path = "./test_corpus_seg/"  # 分词之后的语料库储存路径
    corpus_segment(corpus_path, seg_path, stopwords_file)
    print("测试语料分词结束！！！")